/*
 * pcicfg - PCI drivers & tools for accessing memory location, I/O ports and configuration space.
 * Copyright (C) 2013  Andrew Karasyov
 *
 * This library is free software: you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation, either
 * version 3 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */


#include "pcidevicemanager.h"
#include <fcntl.h>
#include <sys/ioctl.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <iomanip>
#include <string.h>
#include <unistd.h>
#include <assert.h>


PCIDeviceManager::PCIDeviceManager( const char* device_name )
     : PciCfgDeviceHandle( 0 )
{
	// Open driver device
	if( ( PciCfgDeviceHandle = open( device_name, O_RDWR ) ) < 0 )
	{
		std::cout << "Open device error: \"" << device_name << "\"." << std::endl;
		PciCfgDeviceHandle = 0;
		return;
	}
}


PCIDeviceManager::~PCIDeviceManager()
{
	close( PciCfgDeviceHandle );
}


bool PCIDeviceManager::BuildBDFMask( std::string bdf_mask_str, BDFMask& bdf_mask )
{
	if ( bdf_mask_str.empty() )
		return false;

	// The format is [[<bus>]:][<slot>][.[<func>]]. Maximum size is 7.
	if ( bdf_mask_str.length() > 7 )
		return false;

	// Finds ':' and '.' positions
	std::string::size_type colon_pos = bdf_mask_str.find_first_of(':');
	std::string::size_type dot_pos   = bdf_mask_str.find_first_of('.', ((colon_pos != std::string::npos) ? colon_pos + 1 : 0) );
	bool is_colon_present = (colon_pos != std::string::npos);
	bool is_dot_present   = (dot_pos != std::string::npos);

	// If dot is before colon it is an error.
	if( is_colon_present && is_dot_present && colon_pos > dot_pos )
		return false;

	// [TODO]: Put more accurate BDF mask checking.

	unsigned char bus  = 0;
	unsigned char dev  = 0;
	unsigned char func = 0;
	bool isBus         = false;
	bool isDev         = false;
	bool isFunc        = false;

	// First component - bus. It can be only before ':'.
	if ( is_colon_present && (colon_pos > 0) )
	{
		bus = strtol (bdf_mask_str.substr( 0, colon_pos).c_str(), NULL, 16);
		isBus = true;
	}

	// second component - device.
	// It can be :00., :00, 00., 00
	if ( is_colon_present && is_dot_present )
	{
		dev = strtol( bdf_mask_str.substr( colon_pos + 1, dot_pos - colon_pos).c_str(), NULL, 16 );
		isDev = true;
	}
	// If colon is in and more characters in string.
	else if ( is_colon_present && ( bdf_mask_str.length() > (colon_pos + 1) ) )
	{
		dev = strtol( bdf_mask_str.substr( colon_pos + 1).c_str(), NULL, 16 );
		isDev = true;
	}
	// If dot is in and more characters before dot in string.
	else if ( is_dot_present && (dot_pos > 0) )
	{
		dev = strtol( bdf_mask_str.substr( 0, dot_pos).c_str(), NULL, 16 );
		isDev = true;
	}
	// If dot is in and more characters before dot in string.
	else if ( !is_colon_present && !is_dot_present )
	{
		dev = strtol( bdf_mask_str.c_str(), NULL, 16 );
		isDev = true;
	}

	// Third component - function. It can be only after '.'
	if ( is_dot_present && ( bdf_mask_str.length() > (dot_pos + 1) ) )
	{
		func = strtol( bdf_mask_str.substr( dot_pos + 1).c_str(), NULL, 16 );
		isFunc = true;
	}

	bdf_mask.BDF = PCICFG_MAKE_BUSDEVFN( bus, dev, func );
	bdf_mask.BusExist = isBus;
	bdf_mask.DeviceExist = isDev;
	bdf_mask.FuncExist = isFunc;

	return true;
}


bool PCIDeviceManager::BuildRegister( std::string register_str, SRegister& reg )
{
	if ( register_str.empty() )
		return false;

	// The format is <offset>[.(B|W|D)]
	// Finds '.' positions
	std::string::size_type dot_pos = register_str.find_first_of('.');
	bool is_dot_present = (dot_pos != std::string::npos);

	bool is_ok = false;

	// If dot is present and there is text before and after '.'
	if ( is_dot_present && (dot_pos > 0) && (dot_pos < (register_str.length() - 1) ) )
	{
		reg.Offset = strtol( register_str.substr( 0, dot_pos).c_str(), NULL, 0 );
		std::string length_str = register_str.substr( dot_pos + 1 );
		if ( length_str.length() == 1 )
		{
			if (length_str == "B" || length_str == "b")
			{
				reg.Length = VALUE_LENGTH_BYTE;
				is_ok = true;
			}
			else if (length_str == "W" || length_str == "w")
			{
				reg.Length = VALUE_LENGTH_WORD;
				is_ok = true;
			}
			else if (length_str == "D" || length_str == "d")
			{
				reg.Length = VALUE_LENGTH_DWORD;
				is_ok = true;
			}
		}
	}
	else if ( !is_dot_present )
	{
		reg.Offset = strtol( register_str.c_str(), NULL, 0 );
		reg.Length = VALUE_LENGTH_DWORD;
		is_ok = true;
	}

	return is_ok;
}


bool PCIDeviceManager::BuildAssignment( std::string assignment_str, SAssignment& assign )
{
	if ( assignment_str.empty() )
		return false;

	// The format is <offset>[.(B|W|D)]=<value>
	// Finds '=' positions
	std::string::size_type equal_pos = assignment_str.find_first_of('=');
	bool is_equal_present = (equal_pos != std::string::npos);

	// If there is no '=' or there is no text before or after '=' - return false.
	if ( !is_equal_present ||
		 is_equal_present && equal_pos == 0 ||
		 is_equal_present && equal_pos == (assignment_str.length() - 1) )
		return false;

	bool is_ok = true;

	// Split assignment into 2 parts: register and value:
	std::string register_str = assignment_str.substr( 0, equal_pos );
	std::string value_str  = assignment_str.substr( equal_pos + 1 );

	// Builds register
	is_ok = is_ok && BuildRegister( register_str, assign.Register );

	assign.Value = strtol( value_str.c_str(), NULL, 0 );

	return is_ok;
}


void PCIDeviceManager::PrintHelp() const
{
	std::cout <<

	"PCI configuration utility.\n"
	"Usage: pcicfg [<BDF_mask>] [<switches>]\n"
	"\tBDF_mask\t[[<bus>]:][<slot>][.[<func>]] Show only devices in selected slots (hex). \n\n"
	"Switches:\n"
	"\t-l\t Show device list\n"
	"\t-t\t Show bus tree. <BDF_mask> is ignored for bus tree.\n"
	"\t-x[0..4096]\t Show hex-dump of the configuration space. The size is set in bytes.\n"
	"\t-s<reg>=<value>\t Set configuration space register of given device(s). Value is in hex.\n"
	"\t-r<reg>\t Read configuration space register of given device(s).\n"
	"\t<reg>:\t<offset>[.(B|W|D)] Offset is in dec, oct or hex. \n"
	"\t\tB is for byte, \n"
	"\t\tW is for word and \n"
	"\t\tD is for Dword length.\n";
}


void PCIDeviceManager::PrintBDF( unsigned short bdf ) const
{
	std::cout << std::hex << std::setfill('0')
			  << std::setw(2) << PCICFG_GET_BUS( bdf )
			  << ":"
			  << std::setw(2) << PCICFG_GET_DEV( bdf )
			  << "."
			  << PCICFG_GET_FUNC( bdf );
}


bool PCIDeviceManager::ReadAllDevices()
{
	if ( !PciCfgDeviceHandle )
		return false;

	// Gets pci device count.
	int pci_device_count = ioctl( PciCfgDeviceHandle, PCICFG_IOCQDEVICECOUNT );
	if (pci_device_count < 0 )
	{
		std::cout << "PCICFG_IOCQDEVICECOUNT error." << std::endl;
		return false;
	}

	// Gets pci device descriptions.
	PCIDeviceDescription* dev_descriptions = new PCIDeviceDescription[pci_device_count];

	int ret_success = ioctl( PciCfgDeviceHandle, PCICFG_IOCGDEVICELIST, dev_descriptions );
	if (ret_success >= 0)
		std::copy( dev_descriptions, dev_descriptions + pci_device_count, std::back_inserter(PCIDeviceList) );
	else
		std::cout << "Get device description list failed." << std::endl;

	delete [] dev_descriptions;

	return (ret_success >= 0) ? true : false;
}


void PCIDeviceManager::PrintDeviceList( const BDFMask& bdf_mask )
{
	for (size_t i = 0; i < PCIDeviceList.size(); ++i)
	{
		if ( MatchBDF( PCIDeviceList[i].BDF, bdf_mask ) )
		{
			PrintBDF( PCIDeviceList[i].BDF );
			if ( PCIDeviceList[i].IsBridge )
				std::cout << " (bridge to: " << std::hex << std::setfill('0') << std::setw(2) << static_cast<unsigned int>(PCIDeviceList[i].SubordinateBus) << ")";
			std::cout << std::endl;
		}
	}
}


void PCIDeviceManager::PrintDeviceTree()
{
	// Builds device tree.
	DeviceTree device_tree;
	for (size_t i = 0; i < PCIDeviceList.size(); ++i)
	{
		unsigned short curr_dev_bus = PCICFG_GET_BUS( PCIDeviceList[i].BDF );
		DeviceList& device_list = device_tree[curr_dev_bus];
		device_list.push_back( PCIDeviceList[i] );
	}

	// Output device tree.
	PrintDeviceTreeBus(0, device_tree, 0);
	std::cout << std::endl;
}


void PCIDeviceManager::PrintAllDeviceCfgSpace( unsigned int req_cfg_size, const BDFMask& bdf_mask )
{
	req_cfg_size = ( (req_cfg_size > PCICFG_CFGSPACE_LENGTH_MAX_BYTES) ? PCICFG_CFGSPACE_LENGTH_MAX_BYTES : req_cfg_size );

	DeviceListIt it = PCIDeviceList.begin();
	for( ; it != PCIDeviceList.end(); ++it)
	{
		if ( MatchBDF( (*it).BDF, bdf_mask ) )
		{
			PrintBDF( (*it).BDF );

			std::cout << std::endl;

			PrintDeviceCfgSpace( (*it).BDF, req_cfg_size );

			std::cout << std::endl;
		}
	}
}


void PCIDeviceManager::AssignRegisterOnDevices( const BDFMask& bdf_mask, const SAssignment& assign )
{
	DeviceListIt it = PCIDeviceList.begin();
	for( ; it != PCIDeviceList.end(); ++it)
	{
		if ( MatchBDF( (*it).BDF, bdf_mask ) )
			AssignRegister( (*it).BDF, assign );
	}
}


void PCIDeviceManager::ReadRegisterOnDevices( const BDFMask& bdf_mask, const SRegister& reg )
{
	DeviceListIt it = PCIDeviceList.begin();
	for( ; it != PCIDeviceList.end(); ++it)
	{
		if ( MatchBDF( (*it).BDF, bdf_mask ) )
			PrintRegister( (*it).BDF, reg );
	}
}


////////////////////////////////////////////////////////////////////////////////


void PCIDeviceManager::PrintDeviceTreeBus( unsigned short bus, DeviceTree& device_tree, unsigned int depth )
{
	DeviceTree::mapped_type& device_list = device_tree[bus];

	// Generate indent
	std::string indent;
	for (int i = 0; i < depth + 1; ++i)
	{
		if (i > 0)
			indent += "|           ";
		else
			indent += "      ";
	}

	// "-[xx]-" size is 6
	std::cout << "-[" << std::hex << std::setfill('0') << std::setw(2) << bus << "]-";

	DeviceTree::mapped_type::iterator it = device_list.begin();
	for(; it != device_list.end(); ++it )
	{
		if (it != device_list.begin())
			std::cout << indent;

		bool is_last_entry = (it +1) == device_list.end();
		if (is_last_entry)
			std::cout << "\\";
		else
			std::cout << "+";

		std::cout << "-" << std::hex << std::setfill('0') << std::setw(2) << PCICFG_GET_DEV( (*it).BDF ) << "." << PCICFG_GET_FUNC( (*it).BDF );
		if ( (*it).IsBridge )
			PrintDeviceTreeBus( (*it).SubordinateBus, device_tree, depth + 1 );

		if (!is_last_entry)
			std::cout << std::endl;
	}
}


void PCIDeviceManager::PrintDeviceCfgSpace( unsigned short bdf, unsigned int req_cfg_size )
{
	req_cfg_size = ( (req_cfg_size > PCICFG_CFGSPACE_LENGTH_MAX_BYTES) ? PCICFG_CFGSPACE_LENGTH_MAX_BYTES : req_cfg_size );

	// Gets 4K configuration space buffer.
	unsigned char cfg_buffer[PCICFG_CFGSPACE_LENGTH_MAX_BYTES];
	int pcicfg_dev = 0;
	int ret_success = -1;

	memset( cfg_buffer, 0x00, PCICFG_CFGSPACE_LENGTH_MAX_BYTES );
	PCIDeviceValue dev_4k_buf;
	dev_4k_buf.BDF = bdf;
	dev_4k_buf.Buffer = cfg_buffer;

    ret_success = ioctl( PciCfgDeviceHandle, PCICFG_IOCXCFG4K, &dev_4k_buf );
	if (ret_success >= 0)
	{
		// Print out config space
		int offset = 0;
		const unsigned int bytes_per_line = 16;
		for (; offset < req_cfg_size; offset += bytes_per_line)
		{
			std::cout << std::hex << std::setw(3) << std::setfill('0') << offset << ": ";
			int i = 0;
			for (; i < bytes_per_line; ++i)
			{
				std::cout << std::hex << std::setw(2) << std::setfill('0') << (int)(dev_4k_buf.Buffer[offset + i]) << " ";
				if ( !((i+1)%4) && ((i+1)%(bytes_per_line)) )
					std::cout << "| ";
			}

			std::cout << std::endl;
		}
	}
	else
		std::cout << "Get device configuration space failed." << std::endl;
}


void PCIDeviceManager::AssignRegister( unsigned short bdf, const SAssignment& assign )
{
    struct PCIDeviceValue pci_val;
    pci_val.BDF = bdf;
    pci_val.Offset = assign.Register.Offset;

	// Choose command depending on length.
	int command = 0;
	std::string length_str = "<unknown length>";
	switch( assign.Register.Length )
	{
		case VALUE_LENGTH_BYTE:
			command = PCICFG_IOCSCFGBYTE;
			length_str = "byte";
            pci_val.Byte0 = assign.Value;
			break;

		case VALUE_LENGTH_WORD:
			command = PCICFG_IOCSCFGWORD;
			length_str = "word";
            pci_val.Word0 = assign.Value;
            break;

		case VALUE_LENGTH_DWORD:
			command = PCICFG_IOCSCFGDWORD;
			length_str = "dword";
            pci_val.Dword = assign.Value;
            break;

		default:
			assert(false);
	}

	int ret_success = ioctl( PciCfgDeviceHandle, command, &pci_val);
	if (ret_success >= 0)
	{
		PrintBDF( pci_val.BDF );
        std::cout << ": successfully write " << length_str << " # " << (int)(pci_val.Offset) << std::endl;
	}
	else
	{
		PrintBDF( pci_val.BDF );
        std::cout << ": Error writing " << length_str << " # " << (int)(pci_val.Offset) << std::endl;
	}
}


void PCIDeviceManager::PrintRegister( unsigned short bdf, const SRegister& reg )
{
    struct PCIDeviceValue pci_val;
    pci_val.BDF = bdf;
    pci_val.Offset = reg.Offset;

	// Choose command depending on length.
	int command = 0;
	std::string length_str = "<unknown length>";
	switch( reg.Length )
	{
		case VALUE_LENGTH_BYTE:
			command = PCICFG_IOCXCFGBYTE;
			length_str = "byte";
			break;

		case VALUE_LENGTH_WORD:
			command = PCICFG_IOCXCFGWORD;
			length_str = "word";
			break;

		case VALUE_LENGTH_DWORD:
			command = PCICFG_IOCXCFGDWORD;
			length_str = "dword";
			break;

		default:
			assert(false);
	}

	int ret_success = ioctl( PciCfgDeviceHandle, command, &pci_val);
	if (ret_success >= 0)
	{
		PrintBDF( pci_val.BDF );
        std::cout << ": " << length_str << " #" << reg.Offset << " = ";
        switch( reg.Length )
        {
            case VALUE_LENGTH_BYTE:
                std::cout << (int)(pci_val.Byte0);
                break;

            case VALUE_LENGTH_WORD:
                std::cout << (int)(pci_val.Word0);
                break;

            case VALUE_LENGTH_DWORD:
                std::cout << (int)(pci_val.Dword);
                break;

            default:
                assert(false);
        }
        std::cout << std::endl;
	}
	else
	{
		PrintBDF( pci_val.BDF );
		std::cout << ": Error reading " << length_str << " #" << reg.Offset << std::endl;
	}
}


bool PCIDeviceManager::MatchBDF( unsigned short bdf, const BDFMask& bdf_mask )
{
	bool isMatch = true;

	if ( bdf_mask.BusExist )
		isMatch = isMatch && (PCICFG_GET_BUS(bdf) == PCICFG_GET_BUS(bdf_mask.BDF));

	if ( bdf_mask.DeviceExist )
		isMatch = isMatch && (PCICFG_GET_DEV(bdf) == PCICFG_GET_DEV(bdf_mask.BDF));

	if ( bdf_mask.FuncExist )
		isMatch = isMatch && (PCICFG_GET_FUNC(bdf) == PCICFG_GET_FUNC(bdf_mask.BDF));

	return isMatch;
}
